"""
数据库工具模块
"""
import logging
import pandas as pd
import sqlalchemy
from sqlalchemy import create_engine, text
from urllib.parse import quote_plus

logger = logging.getLogger(__name__)

class DBUtils:
    """数据库工具类"""
    
    @staticmethod
    def create_engine_from_config(config):
        """
        根据配置创建数据库引擎
        
        Args:
            config: 数据源配置
        
        Returns:
            sqlalchemy.engine.Engine: 数据库引擎
        """
        db_type = config.get('type')
        host = config.get('host')
        port = config.get('port')
        username = config.get('username')
        password = config.get('password')
        database = config.get('database')
        charset = config.get('charset', 'utf8mb4')
        
        # 对密码进行URL编码
        if password:
            password = quote_plus(password)
        
        connection_string = None
        
        if db_type == 'mysql':
            connection_string = f"mysql+pymysql://{username}:{password}@{host}:{port}/{database}?charset={charset}"
        elif db_type == 'postgresql':
            connection_string = f"postgresql://{username}:{password}@{host}:{port}/{database}"
        elif db_type == 'oracle':
            connection_string = f"oracle+cx_oracle://{username}:{password}@{host}:{port}/{database}"
        elif db_type == 'sqlserver':
            connection_string = f"mssql+pymssql://{username}:{password}@{host}:{port}/{database}"
        elif db_type == 'sqlite':
            connection_string = f"sqlite:///{database}"
        else:
            raise ValueError(f"不支持的数据库类型: {db_type}")
        
        try:
            engine = create_engine(connection_string, echo=False)
            return engine
        except Exception as e:
            logger.error(f"创建数据库引擎失败: {str(e)}")
            raise
    
    @staticmethod
    def test_connection(config):
        """
        测试数据库连接
        
        Args:
            config: 数据源配置
        
        Returns:
            bool: 连接是否成功
        """
        try:
            engine = DBUtils.create_engine_from_config(config)
            with engine.connect() as conn:
                conn.execute(text("SELECT 1"))
            return True
        except Exception as e:
            logger.error(f"测试数据库连接失败: {str(e)}")
            return False
    
    @staticmethod
    def get_table_schema(engine, table_name, schema=None):
        """
        获取表结构
        
        Args:
            engine: 数据库引擎
            table_name: 表名
            schema: 模式名 (可选)
        
        Returns:
            list: 字段信息列表
        """
        try:
            inspector = sqlalchemy.inspect(engine)
            columns = inspector.get_columns(table_name, schema=schema)
            return columns
        except Exception as e:
            logger.error(f"获取表结构失败: {str(e)}")
            raise
    
    @staticmethod
    def get_tables(engine, schema=None):
        """
        获取所有表名
        
        Args:
            engine: 数据库引擎
            schema: 模式名 (可选)
        
        Returns:
            list: 表名列表
        """
        try:
            inspector = sqlalchemy.inspect(engine)
            tables = inspector.get_table_names(schema=schema)
            return tables
        except Exception as e:
            logger.error(f"获取表名列表失败: {str(e)}")
            raise
    
    @staticmethod
    def execute_query(engine, query, params=None):
        """
        执行查询
        
        Args:
            engine: 数据库引擎
            query: SQL查询语句
            params: 查询参数 (可选)
        
        Returns:
            list: 查询结果
        """
        try:
            with engine.connect() as conn:
                result = conn.execute(text(query), params or {})
                return [dict(row) for row in result]
        except Exception as e:
            logger.error(f"执行查询失败: {str(e)}")
            raise
    
    @staticmethod
    def execute_query_to_df(engine, query, params=None):
        """
        执行查询并返回DataFrame
        
        Args:
            engine: 数据库引擎
            query: SQL查询语句
            params: 查询参数 (可选)
        
        Returns:
            pandas.DataFrame: 查询结果
        """
        try:
            return pd.read_sql(text(query), engine, params=params or {})
        except Exception as e:
            logger.error(f"执行查询失败: {str(e)}")
            raise
    
    @staticmethod
    def execute_non_query(engine, query, params=None):
        """
        执行非查询操作
        
        Args:
            engine: 数据库引擎
            query: SQL语句
            params: 查询参数 (可选)
        
        Returns:
            int: 受影响的行数
        """
        try:
            with engine.connect() as conn:
                with conn.begin():
                    result = conn.execute(text(query), params or {})
                    return result.rowcount
        except Exception as e:
            logger.error(f"执行非查询操作失败: {str(e)}")
            raise
    
    @staticmethod
    def bulk_insert(engine, table_name, data, schema=None):
        """
        批量插入数据
        
        Args:
            engine: 数据库引擎
            table_name: 表名
            data: 数据列表
            schema: 模式名 (可选)
        
        Returns:
            int: 插入的行数
        """
        try:
            if not data:
                return 0
            
            metadata = sqlalchemy.MetaData()
            if schema:
                table = sqlalchemy.Table(table_name, metadata, schema=schema, autoload_with=engine)
            else:
                table = sqlalchemy.Table(table_name, metadata, autoload_with=engine)
            
            with engine.connect() as conn:
                with conn.begin():
                    result = conn.execute(table.insert(), data)
                    return result.rowcount
        except Exception as e:
            logger.error(f"批量插入数据失败: {str(e)}")
            raise
    
    @staticmethod
    def df_to_sql(df, engine, table_name, if_exists='replace', index=False, schema=None):
        """
        将DataFrame写入数据库
        
        Args:
            df: pandas.DataFrame
            engine: 数据库引擎
            table_name: 表名
            if_exists: 如果表存在的处理方式 ('fail', 'replace', 'append')
            index: 是否写入索引
            schema: 模式名 (可选)
        
        Returns:
            None
        """
        try:
            df.to_sql(
                name=table_name,
                con=engine,
                if_exists=if_exists,
                index=index,
                schema=schema
            )
        except Exception as e:
            logger.error(f"将DataFrame写入数据库失败: {str(e)}")
            raise
    
    @staticmethod
    def get_primary_key(engine, table_name, schema=None):
        """
        获取表的主键
        
        Args:
            engine: 数据库引擎
            table_name: 表名
            schema: 模式名 (可选)
        
        Returns:
            list: 主键字段列表
        """
        try:
            inspector = sqlalchemy.inspect(engine)
            pk_constraint = inspector.get_pk_constraint(table_name, schema=schema)
            return pk_constraint.get('constrained_columns', [])
        except Exception as e:
            logger.error(f"获取主键失败: {str(e)}")
            raise
    
    @staticmethod
    def get_foreign_keys(engine, table_name, schema=None):
        """
        获取表的外键
        
        Args:
            engine: 数据库引擎
            table_name: 表名
            schema: 模式名 (可选)
        
        Returns:
            list: 外键信息列表
        """
        try:
            inspector = sqlalchemy.inspect(engine)
            foreign_keys = inspector.get_foreign_keys(table_name, schema=schema)
            return foreign_keys
        except Exception as e:
            logger.error(f"获取外键失败: {str(e)}")
            raise
    
    @staticmethod
    def get_indexes(engine, table_name, schema=None):
        """
        获取表的索引
        
        Args:
            engine: 数据库引擎
            table_name: 表名
            schema: 模式名 (可选)
        
        Returns:
            list: 索引信息列表
        """
        try:
            inspector = sqlalchemy.inspect(engine)
            indexes = inspector.get_indexes(table_name, schema=schema)
            return indexes
        except Exception as e:
            logger.error(f"获取索引失败: {str(e)}")
            raise